package com.project.srtp.secure.util;

import com.project.srtp.secure.dto.req.BusinessOperationReqDTO;
import org.apache.tuweni.bytes.Bytes;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.util.encoders.Hex;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.InputStream;
import java.io.OutputStream;
import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.spec.ECGenParameterSpec;
import java.util.Base64;

/**
 * @author 曹行阳
 * @version 0.2
 * @date 2025/6/11
 * CryptoUtil: 对积动算/解密、指纹、承诺、签名功能工具类
 */
@Component
public class CryptoUtil {

    private KeyPair auditKeyPair;
    private BigInteger pedersenG;
    private BigInteger pedersenH;
    private BigInteger pedersenP;

    private static final int AES_KEY_SIZE = 256;
    private static final int GCM_TAG_LENGTH = 128;
    private static final int IV_SIZE = 12;

    private static final String ALGORITHM = "AES";
    private static final String CIPHER_TRANSFORMATION = "AES/GCM/NoPadding";

    @PostConstruct
    public void init() {
        Security.addProvider(new BouncyCastleProvider());

        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("EC", "BC");
            keyGen.initialize(new ECGenParameterSpec("secp256r1"), new SecureRandom());
            this.auditKeyPair = keyGen.generateKeyPair();
        } catch (Exception e) {
            throw new RuntimeException("ECDSA KeyPair generation failed", e);
        }

        // For consistent Pedersen parameters in distributed settings, use fixed values or config
        this.pedersenP = new BigInteger("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F", 16);
        this.pedersenG = BigInteger.valueOf(2);
        this.pedersenH = BigInteger.valueOf(3);
    }

    public static byte[] generateAesKey() throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(AES_KEY_SIZE);
        SecretKey secretKey = keyGen.generateKey();
        return secretKey.getEncoded();
    }

    public static byte[] encrypt(InputStream plainInput, OutputStream cipherOutput, byte[] rawKey) throws Exception {
        byte[] iv = new byte[IV_SIZE];
        SecureRandom random = new SecureRandom();
        random.nextBytes(iv);

        SecretKeySpec keySpec = new SecretKeySpec(rawKey, ALGORITHM);
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
        GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, spec);

        cipherOutput.write(iv);

        byte[] buffer = new byte[4096];
        int len;
        while ((len = plainInput.read(buffer)) != -1) {
            byte[] chunk = cipher.update(buffer, 0, len);
            if (chunk != null) {
                cipherOutput.write(chunk);
            }
        }
        byte[] finalBytes = cipher.doFinal();
        if (finalBytes != null) {
            cipherOutput.write(finalBytes);
        }
        cipherOutput.flush();
        return iv;
    }

    public static void decrypt(InputStream cipherInput, OutputStream plainOutput, byte[] rawKey) throws Exception {
        byte[] iv = new byte[IV_SIZE];
        if (cipherInput.read(iv) != IV_SIZE) {
            throw new IllegalArgumentException("Invalid cipher stream, cannot read IV");
        }

        SecretKeySpec keySpec = new SecretKeySpec(rawKey, ALGORITHM);
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
        GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH, iv);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, spec);

        byte[] buffer = new byte[4096];
        int len;
        while ((len = cipherInput.read(buffer)) != -1) {
            byte[] plainChunk = cipher.update(buffer, 0, len);
            if (plainChunk != null) {
                plainOutput.write(plainChunk);
            }
        }
        byte[] finalBytes = cipher.doFinal();
        if (finalBytes != null) {
            plainOutput.write(finalBytes);
        }
        plainOutput.flush();
    }

    public static SecretKey bytesToKey(byte[] rawKey) {
        return new SecretKeySpec(rawKey, ALGORITHM);
    }

    public String generateNonce(long timestamp, String deviceId, String sessionId) {
        try {
            SecureRandom random = new SecureRandom();
            long salt = random.nextLong();
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            String input = timestamp + deviceId + sessionId + salt;
            byte[] hash = digest.digest(input.getBytes(StandardCharsets.UTF_8));
            return Bytes.wrap(hash).toHexString();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("SHA-256 not available", e);
        }
    }

    public String generateBlake3Fingerprint(BusinessOperationReqDTO opData, String nonce) throws NoSuchAlgorithmException, NoSuchProviderException {
        String concatenatedData = opData.getUserId() + opData.getOpType() + opData.getValue();
        String input = concatenatedData + nonce;
        MessageDigest hasher = MessageDigest.getInstance("BLAKE3", "BC");
        byte[] hash = hasher.digest(input.getBytes(StandardCharsets.UTF_8));
        return Hex.toHexString(hash);
    }

    public String generatePedersenCommitment(String fingerprint) {
        BigInteger f = new BigInteger(fingerprint, 16);
        BigInteger r = new BigInteger(256, new SecureRandom());
        BigInteger c = pedersenG.modPow(f, pedersenP).multiply(pedersenH.modPow(r, pedersenP)).mod(pedersenP);
        return c.toString(16);
    }

    public String generateRangeProof(int value) {
        return "placeholder-bulletproof-for-value-" + value;
    }

    public String generateEcdsaSignature(String commitment, String rangeProof) {
        try {
            Signature ecdsa = Signature.getInstance("SHA256withECDSA", "BC");
            ecdsa.initSign(auditKeyPair.getPrivate());
            String message = commitment + rangeProof;
            ecdsa.update(message.getBytes(StandardCharsets.UTF_8));
            byte[] signatureBytes = ecdsa.sign();
            return Base64.getEncoder().encodeToString(signatureBytes);
        } catch (Exception e) {
            throw new RuntimeException("ECDSA Signature failed", e);
        }
    }

    public boolean verifyEcdsaSignature(String commitment, String rangeProof, String base64Signature) {
        try {
            Signature ecdsa = Signature.getInstance("SHA256withECDSA", "BC");
            ecdsa.initVerify(auditKeyPair.getPublic());
            ecdsa.update((commitment + rangeProof).getBytes(StandardCharsets.UTF_8));
            return ecdsa.verify(Base64.getDecoder().decode(base64Signature));
        } catch (Exception e) {
            return false;
        }
    }
}
